#include "cphotoncache.h"

using namespace std;

template<int n> struct sortingCriterion
{
    photon *photonTab;
    bool operator() (int p1, int p2)
    {
        if (n == kdnode::dimensionX)
        {
            return photonTab[p1].pos.GetX() < photonTab[p2].pos.GetX();
        }

        if (n == kdnode::dimensionY)
        {
            return photonTab[p1].pos.GetY() < photonTab[p2].pos.GetY();
        }

        if (n == kdnode::dimensionZ)
        {
            return photonTab[p1].pos.GetZ() < photonTab[p2].pos.GetZ();
        }
    }
};

void CPhotonCache::balance(vector<int>& tab, int p)
{
    if (tab.size() == 0)
        return;
    vector<int> tabLess;
    vector<int> tabMore;
    if (2 * p + 1 < 100 * int(photonTab.size()))
    {
        kdnode node;
        node.nPhotonMedian = -1;
        if (tab.size() == 1)
        {
            node.nPhotonMedian = tab[0];
        }
        else
        {
            float minX, minY, minZ;
            float maxX, maxY, maxZ;

            minX = minY = minZ = numeric_limits<float>::max();
            maxX = maxY = maxZ = - numeric_limits<float>::max();
            tabLess.reserve(tab.size());
            tabMore.reserve(tab.size());
            for (int i = 0; i < int(tab.size()); ++i)
            {
                minX = min(minX, (float)photonTab[tab[i]].pos.GetX());
                minY = min(minY, (float)photonTab[tab[i]].pos.GetY());
                minZ = min(minZ, (float)photonTab[tab[i]].pos.GetZ());
                maxX = max(maxX, (float)photonTab[tab[i]].pos.GetX());
                maxY = max(maxY, (float)photonTab[tab[i]].pos.GetY());
                maxZ = max(maxZ, (float)photonTab[tab[i]].pos.GetZ());
            }

            if ( maxX - minX >= maxY - minY )
            {
                if (maxX - minX >= maxZ - minZ)
                {
                    node.dimensionPlane = kdnode::dimensionX;
                    sortingCriterion<kdnode::dimensionX> pred = {&photonTab[0]};
                    sort(tab.begin(), tab.end(), pred);
                    node.nPhotonMedian = tab[int(tab.size()) / 2];
                    for (int i = 0; i< int(tab.size()) / 2; ++i)
                    {
                        tabLess.push_back(tab[i]);
                    }
                    for (int i = int(tab.size()) / 2 + 1; i< int(tab.size()); ++i)
                    {
                        tabMore.push_back(tab[i]);
                    }
                }
                else
                {
                    node.dimensionPlane = kdnode::dimensionZ;
                    sortingCriterion<kdnode::dimensionZ> pred = {&photonTab[0]};
                    sort(tab.begin(), tab.end(), pred);
                    node.nPhotonMedian = tab[int(tab.size()) / 2];
                    for (int i = 0; i< int(tab.size()) / 2; ++i)
                    {
                        tabLess.push_back(tab[i]);
                    }
                    for (int i = int(tab.size()) / 2 + 1; i< int(tab.size()); ++i)
                    {
                        tabMore.push_back(tab[i]);
                    }
                }
            }
            else
            {
                if (maxY - minY >= maxZ - minZ)
                {
                    node.dimensionPlane = kdnode::dimensionY;
                    sortingCriterion<kdnode::dimensionY> pred = {&photonTab[0]};
                    sort(tab.begin(), tab.end(), pred);
                    node.nPhotonMedian = tab[int(tab.size()) / 2];
                    for (int i = 0; i< int(tab.size()) / 2; ++i)
                    {
                        tabLess.push_back(tab[i]);
                    }
                    for (int i = int(tab.size()) / 2 + 1; i< int(tab.size()); ++i)
                    {
                        tabMore.push_back(tab[i]);
                    }
                }
                else
                {
                    node.dimensionPlane = kdnode::dimensionZ;
                    sortingCriterion<kdnode::dimensionZ> pred = {&photonTab[0]};
                    sort(tab.begin(), tab.end(), pred);
                    node.nPhotonMedian = tab[int(tab.size()) / 2];
                    for (int i = 0; i< int(tab.size()) / 2; ++i)
                    {
                        tabLess.push_back(tab[i]);
                    }
                    for (int i = int(tab.size()) / 2 + 1; i< int(tab.size()); ++i)
                    {
                        tabMore.push_back(tab[i]);
                    }
                }
            }
        }
        if (p >= int(nodeTab.size()) )
            nodeTab.resize(p + 1);
        nodeTab[p] = node;
    }
    balance(tabLess, 2 * p);
    balance(tabMore, 2 * p + 1);
}
